import os
from PIL import Image, ImageFile
import sys
import io

# 允许处理大图片
ImageFile.LOAD_TRUNCATED_IMAGES = True
Image.MAX_IMAGE_PIXELS = None

def compress_image(image_path, max_size_kb=512):
    """压缩图片到指定大小（KB）以下，保持原始宽高比"""
    # 获取原始图片信息
    file_name, file_ext = os.path.splitext(image_path)
    file_ext = file_ext.lower()
    
    # 打开图片
    img = Image.open(image_path)
    
    # 创建输出文件路径
    output_path = f"{file_name}_compressed{file_ext}"
    
    # 如果文件已经小于目标大小，不做任何修改
    file_size_kb = os.path.getsize(image_path) / 1024
    if file_size_kb <= max_size_kb:
        print(f"{image_path} 已经小于 {max_size_kb}KB，跳过处理")
        return
    
    # 计算目标大小（字节）
    target_size = max_size_kb * 1024
    
    # 对PNG格式进行特殊处理
    if file_ext in ['.png']:
        print(f"对PNG文件 {image_path} 进行特殊压缩处理...")
        
        # 尝试1：先转换为JPG格式并压缩
        print("尝试转换为JPG格式并压缩...")
        jpg_output_path = f"{file_name}_compressed.jpg"
        
        # 处理透明度 - 如果有透明背景，转换为白色背景
        if img.mode == 'RGBA' or 'transparency' in img.info:
            background = Image.new('RGB', img.size, (255, 255, 255))
            if img.mode == 'RGBA':
                background.paste(img, mask=img.split()[3])  # 3 is the alpha channel
            else:
                background.paste(img)
            img_jpg = background
        else:
            img_jpg = img.convert('RGB')
        
        # 使用质量压缩策略处理JPG
        quality = 90  # 初始质量设置
        current_size = float('inf')
        
        while True:
            img_jpg.save(jpg_output_path, format='JPEG', optimize=True, quality=quality)
            current_size = os.path.getsize(jpg_output_path)
            
            if current_size <= target_size or quality <= 10:
                break
            
            # 根据当前大小与目标大小的比例调整质量
            size_ratio = current_size / target_size
            if size_ratio > 2:
                quality -= 15
            elif size_ratio > 1.5:
                quality -= 10
            else:
                quality -= 5
            
            quality = max(10, quality)
        
        print(f"  已转换为JPG格式，大小={current_size/1024:.2f}KB, 质量={quality}")
        
        # 如果JPG转换成功满足大小要求，就使用JPG结果
        if current_size <= target_size:
            # 重命名JPG文件为PNG扩展名（保持原扩展名）
            if os.path.exists(output_path):
                os.remove(output_path)
            output_path_jpg_as_png = f"{file_name}_compressed_png_as_jpg.jpg"
            if os.path.exists(output_path_jpg_as_png):
                os.remove(output_path_jpg_as_png)
            os.rename(jpg_output_path, output_path_jpg_as_png)
            print(f"  注意：图片已转换为JPG格式并保存为{output_path_jpg_as_png}")
        else:
            # 尝试2：如果JPG转换后仍不满足，尝试调整图片尺寸
            print("JPG格式仍不满足大小要求，尝试调整图片尺寸...")
            # 使用更直接的方式计算调整比例
            scale_factor = (target_size / current_size) ** 0.5
            
            # 使用稍微更激进的压缩比例
            scale_factor = min(scale_factor, 0.9)  # 最多缩小到原尺寸的90%
            
            new_width = max(100, int(img.width * scale_factor))  # 确保宽度至少为100
            new_height = max(100, int(img.height * scale_factor))  # 确保高度至少为100
            
            # 调整尺寸
            img_resized = img_jpg.resize((new_width, new_height), Image.LANCZOS)
            img_resized.save(jpg_output_path, format='JPEG', optimize=True, quality=quality)
            current_size = os.path.getsize(jpg_output_path)
            
            print(f"  已调整尺寸：宽={new_width}, 高={new_height}, 大小={current_size/1024:.2f}KB")
            
            # 重命名调整后的JPG文件
            if os.path.exists(output_path):
                os.remove(output_path)
            output_path_jpg_as_png = f"{file_name}_compressed_png_as_jpg.jpg"
            if os.path.exists(output_path_jpg_as_png):
                os.remove(output_path_jpg_as_png)
            os.rename(jpg_output_path, output_path_jpg_as_png)
            print(f"  注意：图片已转换为JPG格式并保存为{output_path_jpg_as_png}")
        
        # 尝试4：如果上述方法仍不满足，尝试将PNG转换为WebP格式
        if current_size > target_size:
            print("尝试转换为WebP格式...")
            # 创建WebP格式的输出路径
            webp_output_path = f"{file_name}_compressed.webp"
            
            # 根据大小调整WebP的质量
            webp_quality = 80
            
            # 迭代压缩WebP格式
            while True:
                img.save(webp_output_path, format='WebP', quality=webp_quality)
                webp_size = os.path.getsize(webp_output_path)
                
                if webp_size <= target_size or webp_quality <= 10:
                    break
                
                webp_quality -= 5
                
            print(f"  已转换为WebP格式，大小={webp_size/1024:.2f}KB, 质量={webp_quality}")
            
            # 如果WebP格式更有效，使用WebP结果
            if webp_size < current_size:
                current_size = webp_size
                # 重命名WebP文件为PNG扩展名（用户可能希望保持原扩展名）
                if os.path.exists(output_path):
                    os.remove(output_path)
                # 创建一个链接或复制文件，但保留原始扩展名
                output_path_with_webp = f"{file_name}_compressed_png.webp"
                if os.path.exists(output_path_with_webp):
                    os.remove(output_path_with_webp)
                os.rename(webp_output_path, output_path_with_webp)
                print(f"  注意：为了达到大小要求，图片已转换为WebP格式并保存为{output_path_with_webp}")
    else:
        # 对于其他格式（如JPEG），使用质量压缩策略
        quality = 90  # 初始质量设置
        
        # 迭代压缩直到达到目标大小
        while True:
            # 保存图片并设置质量
            img.save(output_path, optimize=True, quality=quality)
            
            # 检查文件大小
            current_size = os.path.getsize(output_path)
            
            if current_size <= target_size or quality <= 10:
                break
            
            # 根据当前大小与目标大小的比例调整质量
            size_ratio = current_size / target_size
            if size_ratio > 2:
                quality -= 15
            elif size_ratio > 1.5:
                quality -= 10
            else:
                quality -= 5
            
            # 确保质量不低于10
            quality = max(10, quality)
    
    print(f"{image_path} 压缩完成，原始大小: {file_size_kb:.2f}KB, 压缩后大小: {current_size/1024:.2f}KB")


def compress_all_images_in_directory(directory='.', max_size_kb=512):
    """递归压缩目录及其所有子目录中的图片文件"""
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
    
    # 递归遍历目录及其所有子目录
    for root, dirs, files in os.walk(directory):
        print(f"正在处理目录: {root}")
        
        for file in files:
            # 跳过已压缩的文件
            if "_compressed" in file:
                print(f"跳过已压缩文件: {os.path.join(root, file)}")
                continue
            
            # 检查是否为图片文件
            _, ext = os.path.splitext(file)
            if ext.lower() in image_extensions:
                file_path = os.path.join(root, file)
                try:
                    compress_image(file_path, max_size_kb)
                except Exception as e:
                    print(f"处理 {file_path} 时出错: {e}")

if __name__ == "__main__":
    # 获取命令行参数，如果没有提供则使用当前目录和512KB
    directory = sys.argv[1] if len(sys.argv) > 1 else '.'
    max_size_kb = int(sys.argv[2]) if len(sys.argv) > 2 else 512
    
    print(f"开始压缩 {directory} 目录下的图片到 {max_size_kb}KB 以下...")
    compress_all_images_in_directory(directory, max_size_kb)
    print("所有图片压缩完成！")